package mosdi.fa;

import java.util.Arrays;
import java.util.List;
import java.util.Random;

import mosdi.util.BitArray;

public class MarkovianTextModel extends AutomatonBasedTextModel {
	private int order;
	private int alphabetSize;
	private int stateCount;
	// every state represents one k-mer (where k is the order of the model)
	// e.g. we have alphabetSize^k states, ordered naturally.
	// this table contains for each state the targetState under character 0
	// (other characters go to subsequent states)
	private int[] transitions;
	// corresponding probabilities (for each state and each character the probability
	// is stored).
	private double[][] transitionProbs;
	// each state's equilibrium probability
	private double[] eqDistribution;
	private double[] cumulativeEqDistribution;

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param text Sample text from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, int alphabetSize, int[] text) {
		this(order, alphabetSize, text, 0.0d);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param text Sample text from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, int alphabetSize, int[] text, double pseudoCounts) {
		if (order<1) throw new IllegalArgumentException("Order must be >=1.");
		this.order = order;
		this.alphabetSize = alphabetSize;
		double[] wordFrequencies = new double[(int)Math.pow(alphabetSize, order+1)];
		Arrays.fill(wordFrequencies, pseudoCounts);
		countWords(text, wordFrequencies);
		buildTables(wordFrequencies);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param text Sample text from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, Alphabet alphabet, String text) {
		this(order,alphabet.size(),alphabet.buildIndexArray(text),0.0d);
	}
	
	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param text Sample text from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, Alphabet alphabet, String text, double pseudoCounts) {
		this(order,alphabet.size(),alphabet.buildIndexArray(text),pseudoCounts);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param texts Sample texts from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, int alphabetSize, List<int[]> texts) {
		this(order, alphabetSize, texts, 0.0d);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param texts Sample texts from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, int alphabetSize, List<int[]> texts, double pseudoCounts) {
		if (order<1) throw new IllegalArgumentException("Order must be >=1.");
		this.order = order;
		this.alphabetSize = alphabetSize;
		double[] wordFrequencies = new double[(int)Math.pow(alphabetSize, order+1)];
		Arrays.fill(wordFrequencies, pseudoCounts);
		for (int[] text : texts) {
			countWords(text, wordFrequencies);
		}
		buildTables(wordFrequencies);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param texts Sample texts from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, Alphabet alphabet, List<String> texts) {
		this(order, alphabet, texts, 0.0d);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param texts Sample texts from which the empiric distribution is to be estimated.
	 */
	public MarkovianTextModel(int order, Alphabet alphabet, List<String> texts, double pseudoCounts) {
		if (order<1) throw new IllegalArgumentException("Order must be >=1.");
		this.order = order;
		this.alphabetSize = alphabet.size();
		double[] wordFrequencies = new double[(int)Math.pow(alphabetSize, order+1)];
		Arrays.fill(wordFrequencies, pseudoCounts);
		for (String text : texts) {
			countWords(alphabet.buildIndexArray(text), wordFrequencies);
		}
		buildTables(wordFrequencies);
	}

	/** Constructor.
	 * 
	 * @param order Order of the model (i.e. length of relevant history)
	 * @param wordFrequencies Table of frequencies (in arbitrary units) of (order+1)-grams
	 *                        in lexicographic order.
	 */
	public MarkovianTextModel(int order, int alphabetSize, double[] wordFrequencies) {
		if (order<1) throw new IllegalArgumentException("Order must be >=1.");
		if (wordFrequencies.length!=(int)Math.pow(alphabetSize,(order+1))) throw new IllegalArgumentException("Array wordFrequencies has invalid length.");
		this.order = order;
		this.alphabetSize = alphabetSize;
		buildTables(wordFrequencies);
	}
	
	/** Counts words in s of length this.order+1 and adds them to wordFrequencies. */
	private void countWords(int[] s, double[] wordFrequencies) {
		int wordIndex = 0;
		int h = ((int)Math.pow(alphabetSize, order));
		for (int i=0; i<s.length; ++i) {
			wordIndex = (wordIndex%h)*alphabetSize + s[i];
			if (i>=order) wordFrequencies[wordIndex]+=1;
		}
	}
	
	private void buildTables(double[] wordFrequencies) {
		// Log.println(Log.Level.DEBUG, printWordFrequencies(wordFrequencies));
		stateCount = (int)Math.pow(alphabetSize,order);
		transitions = new int[stateCount];
		transitionProbs = new double[stateCount][];
		int wordIndex = 0;
		double totalSum = 0.0;
		eqDistribution = new double[stateCount];
		int h = ((int)Math.pow(alphabetSize, order-1));
		for (int state=0; state<stateCount; ++state) {
			transitions[state] = (state%h)*alphabetSize;
			double sum = 0.0;
			for (int c=0; c<alphabetSize; ++c) sum+=wordFrequencies[wordIndex+c];
			transitionProbs[state] = new double[alphabetSize];
			for (int c=0; c<alphabetSize; ++c) transitionProbs[state][c]=wordFrequencies[wordIndex++]/sum;
			eqDistribution[state]=sum;
			totalSum+=sum;
		}
		
		
		double tmpProb = 0.0d; 
		for (int i = 0; i < stateCount; ++i) {
			for (int j = 0; j < alphabetSize; ++j) {
				tmpProb = getProbability(i, j);
				if (tmpProb == 0.0d) {
					throw new IllegalArgumentException("The Markovian Modell has a transition (State, character) = ("+i+","+j+") with a value equals 0");
					// throw new RuntimeException("Wahrscheinlichkeit vom
					// Zustand "+i+" mit Zeichen "+alphabet.get(j)+"
					// weiterzukommen ist 0");
				} 
			}
		}
		for (int i=0; i<stateCount; ++i) eqDistribution[i]/=totalSum;
		convergeToEquilibrium(1e-13);
	}
	
	private void convergeToEquilibrium(double accuracy) {
		// Log.println(Log.Level.DEBUG, Arrays.toString(eqDistribution));
		int n = 0;
		final int maxsteps = 10000;
		while (true) {
			double[] p = new double[eqDistribution.length];
			for (int state=0; state<stateCount; ++state) {
				for (int c=0; c<alphabetSize; ++c) {
					p[getTransitionTarget(state,c)]+=eqDistribution[state]*transitionProbs[state][c];
				}
			}
			// check if equilibrium is reached
			double maxdiff = 0.0;
			for (int state=0; state<stateCount; ++state) {
				double diff = Math.abs(eqDistribution[state]-p[state]);
				if (diff>maxdiff) maxdiff=diff;
			}
			eqDistribution=p;
			// Log.printf(Log.Level.DEBUG, "Convergence after step %d: %e%n", n, maxdiff);
			if (maxdiff<=accuracy) break;
			if (n>=maxsteps) throw new IllegalStateException("Markov chain does not converge!");
			n+=1;
		}
		// Log.println(Log.Level.DEBUG, Arrays.toString(eqDistribution));
	}
	
	@Override
	public int getAlphabetSize() {
		return alphabetSize;
	}

	/** Returns the string to which the given state corresponds. (Each state
	 *  corresponds to a string of length "order". If in this state, the returned
	 *  string matches the lastly generated characters. */
	@Override
	public int[] getContext(int state) {
		int[] s = new int[order];
		for (int i=0; i<order; ++i) {
			int c = (state/(int)Math.pow(alphabetSize, order-i-1))%alphabetSize;
			s[i] = c;
		}
		return s;
	}
	
//	public String printWordFrequencies(double[] wordFrequencies) {
//		StringBuffer sb = new StringBuffer();
//		int l = order+1;
//		double sum = 0.0;
//		for (int i=0; i<wordFrequencies.length; ++i) sum+=wordFrequencies[i];
//		for (int i=0; i<wordFrequencies.length; ++i) {
//			for (int j=0; j<l; ++j) {
//				int c = (i/(int)Math.pow(alphabetSize, l-j-1))%alphabetSize;
//				sb.append(alphabet.get(c));
//			}
//			sb.append('\t');
//			sb.append(wordFrequencies[i]);
//			sb.append('\t');
//			sb.append(wordFrequencies[i]/sum);
//			sb.append('\n');
//		}
//		return sb.toString();
//	}
	
	public int getTransitionTarget(int sourceState, int character) {
		return transitions[sourceState] + character;
	}
	
	/** Probability that, in the given state, the given character is generated. */
	public double getProbability(int state, int character) {
		return transitionProbs[state][character];
	}
	
	public int getOrder() { return order; }
	
	/** Given a value between 0.0 and 1.0, find the corresponding state (w.r.t. to
	 *  the cumulative distribution function). */
	private int findState(double p) { 
		if (cumulativeEqDistribution==null) {
			cumulativeEqDistribution = new double[stateCount];
			cumulativeEqDistribution[0] = eqDistribution[0];
			for (int i=1; i<stateCount; ++i) {
				cumulativeEqDistribution[i] = cumulativeEqDistribution[i-1] + eqDistribution[i];
			}
			cumulativeEqDistribution[stateCount-1]=1.0;
		}
		// binary search for the smallest index i, s.t. cumulativeEqDistribution[i]>=p
		if (cumulativeEqDistribution[0]>=p) return 0;
		int l = 0;
		int r = stateCount-1;
		while (true) {
			int k = (l+r+1)/2;
			if (cumulativeEqDistribution[k]<p) l = k;
			else {
				if (cumulativeEqDistribution[k-1]<p) return k; 
				r = k;
			}
		}
	}
	
	public int[] generateRandomText(int length) {
		int[] s = new int[length];
		Random random = new Random();
		// choose start
		int state = findState(random.nextDouble());
		s[0] = state%alphabetSize;
		for (int i=1; i<length; ++i) {
			double p = random.nextDouble();
			int c = 0;
			while ((p>transitionProbs[state][c]) && (c<alphabetSize)) {
				p-=transitionProbs[state][c];
				c+=1;
			}
			state = getTransitionTarget(state, c);
			s[i] = state%alphabetSize;
		}
		return s;
	}
	/**
	 * 
	 * @param pattern a distict pattern represented as an integer-array
	 * @return the Expected occurence of the given pattern in the text represented with this MarkovianTextModel with a text-length 1 
	 */	
	public double calculateExpectation(int[] pattern) {
		double[] stateDist = new double[this.getStateCount()];
		double[] tmp;
		double[] newstateDist = new double[this.getStateCount()];

		for (int i = 0; i < this.getStateCount(); ++i) {
			newstateDist[i] = getEquilibriumProbability(i);
		}

		for (int j = 0; j < pattern.length; ++j) {
			tmp = newstateDist;
			newstateDist = stateDist;
			stateDist = tmp;

			// clean up the array
			Arrays.fill(newstateDist, 0.0);

			for (int i = 0; i < this.getStateCount(); ++i) {
				if (stateDist[i] == 0.0d)
					continue;
				newstateDist[this.getTransitionTarget(i, pattern[j])] += (stateDist[i] * this.getProbability(i, pattern[j]));
			}
		}
		double Exp = 0.0d;
		for (int i = 0; i < newstateDist.length; ++i)
			Exp += newstateDist[i];
		return Exp;
	}

	/**
	 * @param pattern A pattern in the format of GeneralizedString
	 * @return the Expected occurence of the given pattern in the text represented with this MarkovianTextModel with a...
	 */
	public double calculateExpectation(GeneralizedString pattern) {
		double[] newstateDist = new double[this.getStateCount()];
		double[] stateDist = new double[this.getStateCount()];
		double[] tmp;

		System.arraycopy(this.eqDistribution, 0, newstateDist, 0,
				eqDistribution.length);

		for (int j = 0; j < pattern.length(); ++j) {
			tmp = newstateDist;
			newstateDist = stateDist;
			stateDist = tmp;

			// clean up the array
			Arrays.fill(newstateDist, 0.0);
			BitArray bitArray = pattern.getPosition(j);

			for (int i = 0; i < this.getStateCount(); ++i) {
				if (stateDist[i] == 0.0d) continue;
				for (int k = 0; k < bitArray.size(); ++k) {
					// geh über alle möglichen k Buchstaben des Alphabets:
					// prüfe alle mölichen k Buchstaben an der Stelle j im
					// Pattern
					// wenn Bit true ist, ist dieser distinct-Character mit in
					// die Berechnung einzubeziehen

					if (bitArray.get(k)) {
						// newstateDist[this.getTransitionTarget(i, pattern[j])]
						// +=
						// (stateDist[i] * this.getProbability(i, pattern[j]));
						newstateDist[this.getTransitionTarget(i, k)] += (stateDist[i] * this
								.getProbability(i, k));
					}
				}
			}
		}
		double Exp = 0.0d;
		for (int i = 0; i < newstateDist.length; ++i) {
			Exp += newstateDist[i];
		}		
		return Exp;
	}

	public int getStateCount() { return stateCount; }
	
	public double getEquilibriumProbability(int state) { return eqDistribution[state]; }

	/**
	 * @param generalizedAlphabet
	 * @return
	 */
	public double[] buildMinProbabilityTable(BitArray[] generalizedAlphabet) {		
		return buildMinOrMaxProbabilityTable(generalizedAlphabet, false);
	}

	/**
	 * @param generalizedAlphabet
	 * @return
	 */
	public double[] buildMaxProbabilityTable(BitArray[] generalizedAlphabet) {
		return buildMinOrMaxProbabilityTable(generalizedAlphabet, true);
	}
	
	
	/**
	 * @param generalizedAlphabet a valid Alphabet
	 * @param max if true it calculates a table for the MaxProbability.<br> if false it calculates a table for the MinProbability
	 * @return
	 */
	private double[] buildMinOrMaxProbabilityTable(BitArray[] generalizedAlphabet, boolean max) {
		
		double[] minMaxTable = new double[generalizedAlphabet.length]; //declare tabular for min (or max)-values
		
		if(max) {                                                         //collect max-probabilities
			Arrays.fill(minMaxTable, 0.0d);                               //set all values on 0 [because probabilities are higher than 0]
		}
		else {                                                            //collect min-probabilities
			Arrays.fill(minMaxTable, 1.0d);                               //set all value on 1 [because all probabilities are lower than 1]
		}

		double charProb;                                                  //use a variable containing the 'best' probability for a generalised character
		for (int k = 0; k < generalizedAlphabet.length; ++k) {            // for each 'character' or 'character-group'
			for (int i = 0; i < stateCount; ++i) {                        //   for each state (collect min/max) 
				charProb = 0.0d;                                          //     reset Sum for each state
				for (int x = 0; x < generalizedAlphabet[k].size(); ++x) { //     for each realisation of the 'character-group'
					if(generalizedAlphabet[k].get(x)) {                   //       if 'character'-group {k} contains this elemental 'char' {x}
						                                                  //         collect sum of probabilities
						charProb += this.getProbability(i, x);            //         add probability of reading char [x] given the state {i} 
					}
				}
				//got the value
				if(max) {                                                 //     overwrite value if wise 
					minMaxTable[k] = Math.max(minMaxTable[k], charProb);
				}
				else {
					minMaxTable[k] = Math.min(minMaxTable[k], charProb);
				}
			}
		}
		return minMaxTable;
	}
	

	public class ExpectationBoundCalculator {
		/**
		 * position within distribution (with respect to the pattern-length) where the distribution is still valid
		 * validUpTo = 0: position pattern[0] is calculated within distributions
		 * distributions[validUpTo][x] has to be valid
		 */
		private int validUpTo;
		private BitArray[] generalizedAlphabet;
		private int[] pattern;
		private double min,max;
		private double[] minTable;
		private double[] maxTable;
		private double[][] distributions;

		/**
		 * stores the sum of the distribution. 
		 * sum_i=0^t { sum_j=0^stateCount { distribution[i][j] } } = patternExpectation[t]
		 */
		private double[] patternExpectation;

		/**
		 * 
		 * @param generalizedAlphabet
		 * @param pattern Pattern with respect to generalized alphabet.
		 */
		public ExpectationBoundCalculator(BitArray[] generalizedAlphabet, int[] pattern) {
			this.generalizedAlphabet = generalizedAlphabet;
			this.pattern = pattern;
			
			for(int i=0; i<generalizedAlphabet.length; ++i) {//check whether each 'character' of the generalizedAlphabet has the length of the alphabet of the Markovian Modell
				if(generalizedAlphabet[i].size() != alphabetSize) {
					throw new IllegalArgumentException("The Size of the character of generalisedAlphabet["+i+"] is "+generalizedAlphabet[i].size()+" while the size of alphabet (of the MarkovianTextModell) is "+MarkovianTextModel.this.alphabetSize);
				}
			}
			
			this.distributions = new double[pattern.length][eqDistribution.length];
			this.patternExpectation = new double[pattern.length]; 
			
			minTable = buildMinOrMaxProbabilityTable(this.generalizedAlphabet,false);
			maxTable = buildMinOrMaxProbabilityTable(this.generalizedAlphabet, true);
			
			//calculate the min and max-probability: not regarding the type of character
			min = 5.0d;
			max = -5.0d;
			for(int i = 0; i < minTable.length; ++i) {
				min = Math.min(min, minTable[i]);				
			}
			for(int i = 0; i < maxTable.length; ++i) {
				max = Math.max(max, maxTable[i]);
			}
			validUpTo = -1;
		}
		
		/**
		 * @param patternPosition the prefix pattern[0,..., patternPosition] shall be updated within distributions
		 */
		private void updatePrefixDistribution(int patternPosition) {
			double[] lastDistribution = null;
			while (validUpTo < patternPosition) { // be lazy
				if (validUpTo < 0) {
					lastDistribution = eqDistribution;
				} else {
					lastDistribution = distributions[validUpTo];
				}

				Arrays.fill(distributions[validUpTo + 1], 0.0d); // just to be sure delete the contents

				for (int i = 0; i < eqDistribution.length; ++i) { // for each state in the markovian modell
					for (int j = 0; j < alphabetSize; ++j) { // find correct character
						if (generalizedAlphabet[pattern[validUpTo + 1]].get(j)) {
							distributions[validUpTo + 1][getTransitionTarget(i,j)] += lastDistribution[i] * getProbability(i, j);
						}
					}
				}
				
				//updates the patternExpectation-Array
				patternExpectation[validUpTo+1] = 0.0d;
				for(int i = 0; i < MarkovianTextModel.this.stateCount; ++i) {
					patternExpectation[validUpTo+1] += distributions[validUpTo+1][i];
				} //Expectation calculated for 
				
				++validUpTo;
			}
		}
		
		/**
		 * Calculate lower bound for expectation for a pattern composed of the
		 * prefix pattern[0..prefixLength-1] followed by a suffix consisting of the 
		 * characters given by abelianPatternSuffix.
		 * @param prefixLength an index of the range [0,...,pattern.length-1]
		 * @param abelianPatternSuffix Composition of suffix.
		 * @return
		 */
		public double getLowerBoundGivenPrefix(int prefixLength, int[] abelianPatternSuffix) {
			return getLowerOrUpperBoundGivenPrefix(prefixLength, abelianPatternSuffix, false);
		}
		
		/**
		 * Calculate upper bound for expectation for a pattern composed of the
		 * prefix pattern[0..prefixLength-1] followed by a suffix consisting of the 
		 * characters given by abelianPatternSuffix.
		 * @param prefixLength an index of the range [0,...,pattern.length-1]
		 * @param abelianPatternSuffix Composition of suffix.
		 * @return
		 */
		public double getUpperBoundGivenPrefix(int prefixLength, int[] abelianPatternSuffix) {
			return getLowerOrUpperBoundGivenPrefix(prefixLength, abelianPatternSuffix, true);
		}
		
		/** Returns the expectation of the patterns prefix. */
		public double getPrefixExpectation(int prefixLength) {
			updatePrefixDistribution(prefixLength-1);
			return (prefixLength>0)?patternExpectation[prefixLength-1]:1.0;
		}
		
		/**
		 * 
		 * @param prefixLength the length of the prefix of the given pattern within [0,...,pattern.length]
		 * @param abelianPatternSuffix
		 * @param upper if true it calculates the Upper Bound; if false it calculates the Lower Bound
		 * @return
		 */
		private double getLowerOrUpperBoundGivenPrefix(int prefixLength, int[] abelianPatternSuffix, boolean upper) {
			//no if-clause. update is lazy and will only calculate more if necessary
			updatePrefixDistribution(prefixLength-1);
			
			double bound = (prefixLength>0)?patternExpectation[prefixLength-1]:1.0;
			
			for (int i = 0; i < abelianPatternSuffix.length; ++i) { //for every generalized-character
				bound *= Math.pow(upper?maxTable[i]:minTable[i], (double)abelianPatternSuffix[i]);
			}
			return bound;
		}
		
		
		/**
		 * 
		 * @param position Index in pattern[0,...,n-1] where charakter pattern[position] shall be changed
		 * @param newChar the new character which shall replace pattern[position]
		 */
		public void setPatternPosition(int position, int newChar) {
			if(position < 0 || position >= pattern.length) {
				throw new IllegalArgumentException("position ("+position+") is not valid. Choose within the range [0,"+(pattern.length-1)+"]");
			}
			if (pattern[position]==newChar) return;
			validUpTo = Math.min(validUpTo,position-1); 
			pattern[position] = newChar;
		}
		

	}
}
